import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import random
import matplotlib.pyplot as plt
from collections import deque
import pygame
import time

# 颜色定义保持不变

class SnakeEnv:
    def __init__(self, size=20):
        self.size = size
        self.reset()
        
        pygame.init()
        self.cell_size = 20
        self.screen = pygame.display.set_mode((self.size * self.cell_size, self.size * self.cell_size))
        pygame.display.set_caption('Snake AI')
        self.clock = pygame.time.Clock()
        self.font = pygame.font.Font(None, 36)

    def reset(self):
        self.snake = [(self.size//2, self.size//2)]
        self.food = self._place_food()
        self.score = 0
        self.steps = 0
        self.steps_without_food = 0
        return self._get_state()

    # _place_food 方法保持不变

    def _get_state(self):
        state = np.zeros((4, self.size, self.size), dtype=np.float32)
        
        for i, (x, y) in enumerate(self.snake):
            state[0, x, y] = 1 if i < len(self.snake)-1 else 2
        
        state[1, self.food[0], self.food[1]] = 1
        state[2, 0, :] = state[2, -1, :] = state[2, :, 0] = state[2, :, -1] = 1
        
        # 新增一个通道来表示蛇头方向
        head_x, head_y = self.snake[0]
        if len(self.snake) > 1:
            neck_x, neck_y = self.snake[1]
            direction_x, direction_y = head_x - neck_x, head_y - neck_y
            state[3, head_x, head_y] = (direction_x + 1) * 3 + (direction_y + 1)
        
        return state

    def step(self, action):
        self.steps += 1
        self.steps_without_food += 1
        head_x, head_y = self.snake[0]
        
        if action == 0: head_x -= 1
        elif action == 1: head_x += 1
        elif action == 2: head_y -= 1
        elif action == 3: head_y += 1
        
        if (head_x < 0 or head_x >= self.size or head_y < 0 or head_y >= self.size or
            (head_x, head_y) in self.snake):
            return self._get_state(), -1, True
        
        self.snake.insert(0, (head_x, head_y))
        
        if (head_x, head_y) == self.food:
            self.score += 1
            self.food = self._place_food()
            reward = 1
            self.steps_without_food = 0
        else:
            self.snake.pop()
            reward = -0.01  # 减小每步的负奖励
        
        if self.steps_without_food > 100:  # 如果长时间没有吃到食物，给予惩罚
            reward -= 0.5
            self.steps_without_food = 0
        
        if len(self.snake) == self.size * self.size:
            return self._get_state(), 10, True
        
        return self._get_state(), reward, False

    def render(self):
        self.screen.fill(BLACK)
        
        for segment in self.snake:
            pygame.draw.rect(self.screen, GREEN, (segment[1]*self.cell_size, segment[0]*self.cell_size, self.cell_size, self.cell_size))
        
        pygame.draw.rect(self.screen, RED, (self.food[1]*self.cell_size, self.food[0]*self.cell_size, self.cell_size, self.cell_size))
        
        score_text = self.font.render(f'Score: {self.score}', True, WHITE)
        self.screen.blit(score_text, (10, 10))
        
        pygame.display.flip()
        self.clock.tick(20)  # 稍微加快游戏速度

class DQN(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(DQN, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU()
        )
        
        conv_out_size = self._get_conv_out(input_shape)
        self.fc = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions)
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        conv_out = self.conv(x).view(x.size()[0], -1)
        return self.fc(conv_out)

class DQNAgent:
    def __init__(self, state_shape, n_actions):
        self.state_shape = state_shape
        self.n_actions = n_actions
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        self.policy_net = DQN(state_shape, n_actions).to(self.device)
        self.target_net = DQN(state_shape, n_actions).to(self.device)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        
        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=0.0001)
        self.memory = deque(maxlen=100000)
        
        self.epsilon = 1.0
        self.epsilon_decay = 0.9995  # 降低衰减速度
        self.epsilon_min = 0.01  # 提高最小探索率
        self.gamma = 0.99
        self.batch_size = 64  # 增大批量大小
        self.update_target_steps = 1000  # 增加目标网络更新频率
        self.steps_done = 0

    def select_action(self, state):
        if random.random() > self.epsilon:
            with torch.no_grad():
                state = torch.FloatTensor(state).unsqueeze(0).to(self.device)
                q_values = self.policy_net(state)
                return q_values.max(1)[1].item()
        else:
            return random.randrange(self.n_actions)

    def store_transition(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def update_model(self):
        if len(self.memory) < self.batch_size:
            return
        
        batch = random.sample(self.memory, self.batch_size)
        state_batch, action_batch, reward_batch, next_state_batch, done_batch = zip(*batch)
        
        state_batch = torch.FloatTensor(state_batch).to(self.device)
        action_batch = torch.LongTensor(action_batch).unsqueeze(1).to(self.device)
        reward_batch = torch.FloatTensor(reward_batch).to(self.device)
        next_state_batch = torch.FloatTensor(next_state_batch).to(self.device)
        done_batch = torch.FloatTensor(done_batch).to(self.device)
        
        q_values = self.policy_net(state_batch).gather(1, action_batch)
        next_q_values = self.target_net(next_state_batch).max(1)[0].detach()
        expected_q_values = reward_batch + (1 - done_batch) * self.gamma * next_q_values
        
        loss = F.smooth_l1_loss(q_values, expected_q_values.unsqueeze(1))
        
        self.optimizer.zero_grad()
        loss.backward()
        for param in self.policy_net.parameters():
            param.grad.data.clamp_(-1, 1)
        self.optimizer.step()
        
        self.steps_done += 1
        if self.steps_done % self.update_target_steps == 0:
            self.target_net.load_state_dict(self.policy_net.state_dict())
        
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)

def plot_progress(scores, epsilons):
    plt.figure(figsize=(12, 5))
    plt.subplot(121)
    plt.plot(scores)
    plt.title('Score')
    plt.subplot(122)
    plt.plot(epsilons)
    plt.title('Epsilon')
    plt.show()

def train(env, agent, n_episodes=10000):
    scores = []
    epsilons = []
    
    for episode in range(n_episodes):
        state = env.reset()
        total_reward = 0
        done = False
        
        while not done:
            action = agent.select_action(state)
            next_state, reward, done = env.step(action)
            agent.store_transition(state, action, reward, next_state, done)
            agent.update_model()
            
            state = next_state
            total_reward += reward
            
            env.render()
            
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    return
        
        scores.append(total_reward)
        epsilons.append(agent.epsilon)
        
        if episode % 100 == 0:
            print(f"Episode {episode}, Score: {env.score}, Epsilon: {agent.epsilon:.2f}")
            plot_progress(scores, epsilons)

if __name__ == "__main__":
    env = SnakeEnv()
    state_shape = env._get_state().shape
    n_actions = 4
    agent = DQNAgent(state_shape, n_actions)
    train(env, agent)